import tkinter as tk
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from calculation_gui import CalculationWellGUI

class PlotWellGUI(CalculationWellGUI):
    def __init__(self, root, calculator):
        super().__init__(root, calculator)
        
    def plot_results(self, top_elevation, bottom_elevation, ground_elevation, segments, sections):
        # 清空当前图形
        self.ax.clear()
        
        # 设置图形属性
        self.ax.set_title("沉井标高示意图")
        self.ax.set_xlabel("水平位置")
        self.ax.set_ylabel("标高 (m)")
        
        # 计算绘图所需的最大半径
        max_radius = 0
        for seg in segments:
            outer_diameter = seg["outer_diameter"] if "outer_diameter" in seg else 0
            max_radius = max(max_radius, outer_diameter / 2)
        
        # 绘制井体分段
        current_elevation = bottom_elevation
        for i, seg in enumerate(segments):
            outer_diameter = seg["outer_diameter"] if "outer_diameter" in seg else 0
            outer_radius = outer_diameter / 2
            
            # 绘制分段圆柱
            self.ax.fill_between([-outer_radius, outer_radius], current_elevation, 
                                current_elevation + seg["length"], alpha=0.5, 
                                label=f"第{i+1}段" if i == 0 else "")
            
            # 绘制分段顶部边界
            self.ax.plot([-outer_radius, outer_radius], [current_elevation + seg["length"]] * 2, 
                        'k--', linewidth=0.7)
            
            # 添加分段标签
            self.ax.text(outer_radius + 1, current_elevation + seg["length"] / 2, 
                        f"第{i+1}段", ha='left', va='center')
            
            current_elevation = current_elevation + seg["length"]
        
        # 绘制分节边界
        current_elevation = bottom_elevation
        for i, sec in enumerate(sections):
            # 绘制分节顶部边界
            self.ax.plot([-max_radius-1, max_radius+1], [current_elevation + sec["length"]] * 2,
                        'r-', linewidth=1.5, linestyle='--')
            
            # 添加分节标签
            self.ax.text(-max_radius - 2, current_elevation + sec["length"] / 2, 
                        f"第{i+1}节", ha='right', va='center')
            
            current_elevation = current_elevation + sec["length"]
        
        # 绘制井底和井顶
        self.ax.plot([-max_radius-1, max_radius+1], [bottom_elevation] * 2, 'k-', linewidth=1.5)
        self.ax.plot([-max_radius-1, max_radius+1], [top_elevation] * 2, 'k-', linewidth=1.5)
        
        # 绘制设计地面标高
        self.ax.plot([-max_radius-1, max_radius+1], [ground_elevation] * 2, 'g', linewidth=1.5, linestyle='--')
        
        # 添加标高标签
        self.ax.text(max_radius + 1, bottom_elevation, f"井底标高: {bottom_elevation:.2f}m", 
                    ha='left', va='center')
        self.ax.text(max_radius + 1, top_elevation, f"井顶标高: {top_elevation:.2f}m", 
                    ha='left', va='center')
        self.ax.text(max_radius + 1, ground_elevation, f"设计地面标高: {ground_elevation:.2f}m", 
                    ha='left', va='center')
        
        # 设置坐标轴范围
        self.ax.set_xlim([-max_radius - 3, max_radius + 3])
        self.ax.set_ylim([bottom_elevation - 2, top_elevation + 2])
        
        # 添加网格
        self.ax.grid(True, linestyle='--', alpha=0.7)
        
        # 添加图例
        self.ax.legend(loc='upper right')
        
        # 更新画布
        self.canvas.draw()